#include <src/pcps.h>
#include <src/blas_lapack.h>
#include <src/random.h>
#include <src/matrix_math.h>

template <class T>
void test_determinant_lemma_derivative(int argc, char ** argv) {

  std::cout << std::endl;

  struct RandomVec {
    static std::vector<T> get(const int n, const double alpha) {
      std::vector<T> retval(n, pcps::zero<T>());
      for (int i = 0; i < n; i++)
        retval[i] = alpha * pcps::random_number<T>();
      return retval;
    }
  };

  // get matrix dimension
  const int n = std::atoi(argv[1]);

  // create work space
  std::vector<T> work(n*n, pcps::zero<T>());
  std::vector<int> iwork(2*n, 0);

  // build matrix
  std::vector<T> A = RandomVec::get(n*n, 0.1);
  for (int i = 0; i < n; i++)
    A[i*n+i] = 1.0 * pcps::unity<T>() + 0.2 * pcps::random_number<T>();

  // print matrix
  std::cout << "Matrix:" << std::endl;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      std::cout << "   " << pcps::format_number("%8.4f", A[i*n+j]);
    std::cout << std::endl;
  }
  std::cout << std::endl;

  // compute inverse and determinant
  std::vector<T> B(A);
  T orig_det;
  pcps::matrix_inverse_lu(n, orig_det, &B.at(0), &work.at(0), &iwork.at(0));
  
  // print inverse
  std::cout << "Inverse:" << std::endl;
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++)
      std::cout << "   " << pcps::format_number("%8.4f", B[i*n+j]);
    std::cout << std::endl;
  }
  std::cout << std::endl;

  // prepare U and T vectors
  std::vector<T> u1 = RandomVec::get(n, 1.0);
  std::vector<T> u2 = RandomVec::get(n, 1.0);
  std::vector<T> u3 = RandomVec::get(n, 1.0);
  std::vector<T> v1 = RandomVec::get(n, 1.0);
  std::vector<T> v2 = RandomVec::get(n, 1.0);
  std::vector<T> v3 = RandomVec::get(n, 1.0);

  // prepare products of inverse times vectors
  std::vector<T> Bu1(n, pcps::zero<T>());
  std::vector<T> Bu2(n, pcps::zero<T>());
  std::vector<T> Bu3(n, pcps::zero<T>());
  std::vector<T> v1B(n, pcps::zero<T>());
  std::vector<T> v2B(n, pcps::zero<T>());
  std::vector<T> v3B(n, pcps::zero<T>());
  for (int i = 0; i < n; i++)
  for (int j = 0; j < n; j++) {
    Bu1[i] += B[i*n+j] * u1[j];
    Bu2[i] += B[i*n+j] * u2[j];
    Bu3[i] += B[i*n+j] * u3[j];
    v1B[j] += v1[i] * B[i*n+j];
    v2B[j] += v2[i] * B[i*n+j];
    v3B[j] += v3[i] * B[i*n+j];
  }

  struct InnerProduct {
    static T get(const int n, const std::vector<T> & left, const std::vector<T> & right) {
      T retval = pcps::zero<T>();
      for (int i = 0; i < n; i++)
        retval += left[i] * right[i];
      return retval;
    }
  };

  // prepare inner products
  const T ip11 = InnerProduct::get(n, v1, Bu1);
  const T ip12 = InnerProduct::get(n, v1, Bu2);
  const T ip13 = InnerProduct::get(n, v1, Bu3);
  const T ip21 = InnerProduct::get(n, v2, Bu1);
  const T ip22 = InnerProduct::get(n, v2, Bu2);
  const T ip23 = InnerProduct::get(n, v2, Bu3);
  const T ip31 = InnerProduct::get(n, v3, Bu1);
  const T ip32 = InnerProduct::get(n, v3, Bu2);
  const T ip33 = InnerProduct::get(n, v3, Bu3);

  struct ShadedGroupOfTerms {
    static T get(const int di, const int dj, 
                 const std::vector<T> & v1B,
                 const std::vector<T> & Bu1,
                 const std::vector<T> & Bu2,
                 const std::vector<T> & Bu3,
                 const T ip11,
                 const T ip12,
                 const T ip13,
                 const T ip21,
                 const T ip22,
                 const T ip23,
                 const T ip31,
                 const T ip32,
                 const T ip33)
    {
      return -v1B[di] * (   Bu1[dj] * ( pcps::unity<T>() + ip22 + ip33 + ip22*ip33 - ip32*ip23 )
                          + Bu2[dj] * (                  - ip21        + ip23*ip31 - ip21*ip33 )
                          + Bu3[dj] * (                  - ip31        + ip21*ip32 - ip22*ip31 ) );
    }

  };

  // choose element with respect to which we will differentiate
  const int di = pcps::global_lcg() % n;
  const int dj = pcps::global_lcg() % n;

  // get contributions from shaded groups of terms
  T analyt_der = pcps::zero<T>();
  analyt_der += ShadedGroupOfTerms::get(di, dj, v1B, Bu1, Bu2, Bu3, ip11,ip12,ip13,ip21,ip22,ip23,ip31,ip32,ip33);
  analyt_der += ShadedGroupOfTerms::get(di, dj, v2B, Bu2, Bu1, Bu3, ip22,ip21,ip23,ip12,ip11,ip13,ip32,ip31,ip33);
  analyt_der += ShadedGroupOfTerms::get(di, dj, v3B, Bu3, Bu2, Bu1, ip33,ip32,ip31,ip23,ip22,ip21,ip13,ip12,ip11);

  // remember inverse element
  const T inv_element = B[dj*n+di];

  // explicitly construct the new matrix and compute its determinant
  for (int i = 0; i < n; i++)
  for (int j = 0; j < n; j++)
    A[i*n+j] += u1[i]*v1[j] + u2[i]*v2[j] + u3[i]*v3[j];
  B = A;
  T old_det;
  pcps::matrix_inverse_lu(n, old_det, &B.at(0), &work.at(0), &iwork.at(0));

  // get contribution from non-shaded term
  analyt_der += inv_element * old_det / orig_det;

  // print analytic derivative
  std::cout << boost::format("Analytic derivative ratio w.r.t. element (%i,%i):") % di % dj << std::endl;
  std::cout << pcps::format_number("%14.6f", analyt_der) << std::endl;
  std::cout << std::endl;

  // compute numerical derivative
  const T epsilon = 1.0e-6 * pcps::unity<T>();
  A[di*n+dj] += epsilon;
  B = A;
  T new_det;
  pcps::matrix_inverse_lu(n, new_det, &B.at(0), &work.at(0), &iwork.at(0));
  const T num_deriv_ratio = ( new_det - old_det ) / epsilon / orig_det;

  // print the numerical derivative
  std::cout << boost::format("Numeric derivative ratio w.r.t. element (%i,%i):") % di % dj << std::endl;
  std::cout << pcps::format_number("%14.6f", num_deriv_ratio) << std::endl;
  std::cout << std::endl;

  // print the difference between the two
  std::cout << boost::format("Difference:") << std::endl;
  std::cout << pcps::format_number("%14.6f", analyt_der - num_deriv_ratio) << std::endl;
  std::cout << std::endl;

}

template void test_determinant_lemma_derivative<double>(int, char **);
template void test_determinant_lemma_derivative<std::complex<double> >(int, char **);

